Repository navigation
[ATOM] accelerate lm cache: load the CPU prefix when the Qwen3.5 HBM hit is off the chunk grid - #2562
Conversation
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
ZLkanyo009
changed the title
ZLkanyo009
changed the title
…the chunk grid An HBM prefix hit that is not a multiple of the LMCache chunk used to refuse the whole CPU prefix (hbm_off_chunk_grid), even when the CPU tier held a state rung above it. Start the KV leg on the chunk floor instead, reload that partial chunk into private blocks so no shared block is written, and keep the rung above the real HBM hit. If the pool cannot back the private blocks, drop the joint load and keep the HBM hit.
9433d1a to
711ccf5
Compare
yhl-amd
left a comment
yhl-amd
left a comment
There was a problem hiding this comment.
Thanks for this. The direction makes sense, and the scope is limited to the qwen/kimi_k3 offload layouts. Kimi's recipe has hash block == LMCache chunk, so in practice only Qwen3-Next/3.5 change behaviour. A few things I'd like addressed before merge (I only read the code, I didn't run the tests):
- Pool check in
allocate(inline, L1111). This can raiseAssertionErrorfrom_fresh_blockwhen the pool fits exactly. It can also evict the joint's own HBM checkpoint before that checkpoint is pinned. - The unaligned branch drops the compressed-hit claim (inline, L841). This refetches and duplicates KV that is still resident, and in that case the "≤255 tokens / ≤15 blocks" cost bound in the description does not hold.
- Failure exposure. Unaligned hits used to keep their HBM hit. Now they take the joint path, where a failure on either leg disowns to token 0. With a large hit and a nearby rung, one failure costs much more than one success saves (e.g. recomputing 200k tokens vs saving 8k). Could you report the joint-load failure rate against
boundaries, so we can judge the trade? - Stale invariants. These comments no longer hold after this change: in
OffloadJoint(sequence.py), L222 ("num_cached_tokensstays at the HBM prefix…") and L232 (claim_tokens"Abovenum_cached_tokens, not below it"); and the K3 connector's comment that the blocks belowstartare shared (kimi_k3/connector.pyL811). - Test coverage (inline on the new test).
Not a blocker, but on design: a possibly cleaner alternative is to keep the shared claim up to hbm and still retrieve from the chunk floor, but have the block GPU connector drop the leading [floor, hbm) blocks on unpack (map them to a scratch block or skip them). num_cached_tokens and the prefix accounting then stay unchanged, and the private-block path in allocate goes away.
Keep a compressed-hit chunk floor when it is already a legal start, and skip an unaligned joint before it is recorded when the in-use tail does not fit without spending the checkpoint about to be pinned.
|
Addressed the review on
The connector-side alternative (keep |
|
Re-ran the same protocols on GSM8K (TP8, 3-shot, concurrency 65): flexible-extract 0.8749 ± 0.0091 (same as the previous run), strict-match 0.8590 ± 0.0096 (0.0023 lower, inside one standard error). c56, rank-0 lookup, 3600 s: 4093 successful requests, 1 connection reset, 14 cancelled when the 30 s grace expired. Prompt cache hit 94.51% (was 94.33%). Output throughput 1009.0 tok/s (252.2 per GPU; was 958.7). TTFT p50/p99 1219/7343 ms (was 1373/9544). ITL p50 45.4 ms (was 49.0). The request mix of a fixed-duration replay differs across runs, so the small gains are not attributed to the review fixes alone. Nothing is worse than the previous column. Details are in the PR description under "Re-run after the review fixes". |
yhl-amd
left a comment
yhl-amd
left a comment
There was a problem hiding this comment.
Thanks, all five points from my previous review are addressed on 05e502613:
- The pool check moved into
_fit_unaligned_jointbefore the commit. I checked the arithmetic: the re-checked demand,n - in_use([0, claim)), matches whatallocateactually takes.allocateno longer evicts before the restore is pinned. - When the compressed hit's chunk floor is already at or above the gated hit, it is kept as the claim. The "≤255 tokens / ≤15 blocks" bound now holds.
- The failure-rate data answers the exposure question.
- The comments in
OffloadJointand the K3 connector now match the behaviour. - The four new tests cover the cases I asked for.
I also traced the KV leg's start for both the compressed-floor case and the lowered-num_cached_tokens case, and the DecodeScheduler.allocate_waiting path. Each one writes only fresh blocks. CI is green on this commit.
Non-blocking follow-ups, fine as separate PRs:
- This was already there before this PR and is only partly covered here.
can_allocate's fit gate protects the HBM-hit checkpoint, but astate_from_hbmjoint pins the boundary checkpoint inallocate. If the hit image is already pinned by another restore, the gate over-counts by one image (~units_per_checkpointunits), so_fresh_blockcan still raise on a tight pool._fit_unaligned_jointuses the right protected hash, but it only runs when an in-use tail adds demand. Re-checking wheneverdecision.state_from_hbmwould close this, for aligned joints too. Rare in practice:state_hbmwas 0 in both runs. _fit_unaligned_jointre-walks the whole hit to recountnum_new_blocks, duplicating the loop incan_allocate. It only runs when the tail adds demand, so the cost is small. Carrying the count forward would remove the duplicate loop.- With the default rank-0 lookup, the joint-load failure rate is 16%, against 6.5% with all-rank lookup. Your analysis makes a good case for changing the
lookup_server_worker_idsdefault in a separate PR.
LGTM.
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
Reviewed at 05e5026. The main point is reproduced in a unit test. Under an exact-fit pool, an unaligned joint can now raise in allocate, which takes down the scheduler, where the merge base refused it cleanly.
- The re-check misses the pinned checkpoint.
_fit_unaligned_jointre-checks the pool only when the tail adds in-use demand (extra > 0), and returns at once whenclaim >= hbm. Otherwise it keepscan_allocate's check, which protected the HBM-hit checkpoint.allocate, however, pins the boundary checkpoint and only then takes fresh blocks. - The aligned case has the same root problem on the merge base. The fix is narrower than the bug.
- An unaligned admission is now all-or-nothing. Any failure on either leg recomputes from token 0 instead of resuming at the HBM hit.
Repro: block 4, chunk 16. The hit rung (5 blocks) is only in the CPU tier and the boundary rung (8) is in HBM. The tail block [16, 20) is cached and free, and free = demand − units_per_checkpoint. can_allocate returns 5 and commits claim 16 / boundary 32. Then allocate(second, 5) pins checkpoint 8 and _fresh_block raises AssertionError: No PAGE unit available for a fresh KV block. On 6e9b3e1 the same admission is refused as hbm_off_chunk_grid and resumes at 20. With an aligned 4-block hit, 6e9b3e1 crashes the same way. The test is built on TestPagedCopyCheckpoint in tests/test_state_checkpoint.py; I can push it if useful.
Not inline, because these lines are not in the diff:
- Where the fix belongs. The root cause is that, for any joint whose state comes from HBM,
can_allocate'sprotected_hash(block_hashes[num_cached_blocks - 1]) is not the checkpointallocatepins. Re-checking everystate_from_hbmjoint inrecord_allocationagainstprotected=boundary_hash, with demandn − used[0, min(claim, hbm)), would fix the aligned case too and replace the special-case helper with itsextra == 0exit. - A docstring that is no longer true.
_fresh_block's docstring says its raise is unreachable throughSchedulerbecause "allocate protects the one checkpoint it is about to pin". After this PR that is reachable for every unaligned HBM hit. - Metrics.
num_cached_tokensis now the chunk floor, soprefix_cache_hit_tokensand the GPU/CPU tier split count the HBM-resident partial chunk[floor, hbm)as a CPU-tier load (scheduler.py around 2315). The PR's Before vs This-PR tier comparison is therefore not measured the same way.
Details inline.
| block_id = self.kv.lookup(block_hashes[i]) | ||
| if block_id != -1 and self.kv.is_used(block_id): | ||
| extra += 1 | ||
| if extra == 0: |
There was a problem hiding this comment.
These two early returns (claim_tokens >= hbm_tokens above, extra == 0 here) keep can_allocate's check, which protected block_hashes[num_cached_blocks - 1]. But when state_from_hbm, allocate pins boundary_hash (begin_restore) before _fresh_block, and that image's units were counted as spendable. On an exact-fit pool _fresh_block then raises (repro in the review summary). The "compressed floor above the unaligned gated hit" path (claim >= hbm) skips the re-check the same way. Re-checking whenever state_from_hbm, with protected=boundary_hash, regardless of extra, closes it.
There was a problem hiding this comment.
Fixed in 4124e7e: _fit_joint now re-checks every state_from_hbm joint, aligned or not and whether or not the tail adds in-use blocks, against protected=boundary_hash. Your exact-fit repro now skips the joint and keeps the HBM hit instead of raising in _fresh_block (covered in tests/test_state_checkpoint.py).
| num_new_blocks -= 1 | ||
| if decision.state_from_hbm: | ||
| protected: int | None = decision.boundary_hash | ||
| elif num_cached_blocks: |
There was a problem hiding this comment.
When the state comes from the tier, this protects the HBM-hit checkpoint, but in a joint allocate attaches boundary_hash (a tier load) and never pins the hit checkpoint. So units that could really be freed are excluded. Case: the tail is used by a live seq, the tier holds the rung, and the only evictable unit is the hit's own HBM checkpoint. The real demand fits, but this reports no room, so the admission is refused as partial_chunk_no_room and the CPU prefix this PR recovers is lost. The docstring's "the checkpoint allocate is about to pin" does not hold for this branch; protected=None fits it.
There was a problem hiding this comment.
Fixed in 4124e7e: for a tier joint the re-check uses protected=None, so the hit image can fund the private tail and a joint that fits is no longer refused as partial_chunk_no_room.
| extra += 1 | ||
| if extra == 0: | ||
| return decision | ||
| num_new_blocks = self._n_hash_blocks(seq) |
There was a problem hiding this comment.
This rebuilds the num_new_blocks that can_allocate just computed by re-walking [0, num_cached_blocks) with kv.lookup and is_used, after walking [claim, hbm) above. For a 200k-token prompt with 16-token hash blocks that is about 12.5k extra dict lookups per qualifying admission, in the scheduler's admission loop. It also leaves out can_allocate's token_ids collision check, and the two copies of the demand formula can drift (only one guards block_id != -1). Carrying the value on the probe or on JointBoundaryDecision would avoid both.
There was a problem hiding this comment.
Fixed in 4124e7e: the demand is now the num_new_blocks that can_allocate already computed (kept on the seq as _admission_num_new_blocks) plus the in-use blocks of the partial chunk [claim, hbm); the hit is no longer walked a second time.
| assert ( | ||
| claim_tokens % hbs == 0 | ||
| ), "a joint claim below the HBM hit is hash-block aligned" | ||
| num_cached_blocks = claim_tokens // hbs |
There was a problem hiding this comment.
Lowering num_cached_blocks to the chunk floor makes an unaligned admission all-or-nothing. Before, it resumed for certain at the HBM hit from that hit's checkpoint. Now any failure after the joint is committed goes through disown_claimed_prefix and recomputes from token 0, not from hbm. Ways it can fail: the connector refuses in _decide_load_after_alloc (joint_boundary_above_lookup); begin_restore misses because another seq's _fresh_block in the same pass evicted the boundary image; LMCache partly misses (the PR's own data has 16/99 boundaries failing with rank-0 lookup); a state load fails. That can mean recomputing hundreds of thousands of tokens. Falling back to hbm on these paths would keep the old floor; none of them is tested.
There was a problem hiding this comment.
Fixed in 4124e7e: hbm_hit_tokens records the gated hit, and if either leg fails retreat_joint_to_hbm restores the shared blocks and that hit's state, so the request resumes at the HBM hit; token 0 is used only when the hit itself is gone. 0f805f0 fixes the follow-up assert so a retreated joint continues prefill on its existing blocks.
| # refills `[floor, hbm)` privately. | ||
| claim_tokens = ((claim_blocks * hbs) // chunk) * chunk | ||
| claim_tokens = max(claim_tokens, hbm_boundary * hbs) | ||
| if claim_tokens < hbm_tokens: |
There was a problem hiding this comment.
This branch never changes the value. claim_blocks = min(len(block_hashes), candidate) >= hbm_boundary, so the floored claim_tokens is already >= the chunk floor of hbm_tokens; when it is below hbm_tokens it already equals (hbm_tokens // chunk) * chunk, and an aligned hbm_tokens never enters. The real change is dropping max(claim_tokens, hbm_tokens). The conditional and its comment describe a step that never happens; removing the branch says what the code does.
There was a problem hiding this comment.
Removed in 4124e7e. The claim is the compressed hit's chunk floor, and the comment there now says why it is not lifted to hbm_tokens.
…e HBM hit A state-from-HBM joint pinned the boundary image after can_allocate had counted it as free, so an exact-fit pool raised in _fresh_block. Re-check every such joint against that image, and when an unaligned joint fails resume at the gated HBM hit instead of token 0.
|
Addressed the review on
|
Keep the unaligned joint fit check beside the new prefix walk. The walk still supplies the admission's num_new_blocks, and allocate claims only up to the chunk floor.
A failed unaligned joint that recovered the HBM hit kept its block table but cleared both offload-resume flags, so the next admission called allocate and asserted. Mark that path as a failed load so prefill continues from the hit.
|
Crash. On Cause. What "KV load failed" is. The log line is the joint pair, not a KV miss. Over the hour below, that line fired 40 times and Accuracy, measured on
flexible-extract is 0.0061 lower, inside 1 SE of the previous run. Performance, measured on
The fixed-duration replay mix is not the same across runs, so the throughput gap is not only this change. This is LMCache-on only; the on/off comparison in the issue checklist was not re-run. |
|
Accuracy re-run on
flexible-extract is 0.0099 above |
No conflicts. Picks up #2669 (fix(mla): migrate Triton decode to maintained aiter API): the 10-20 heavy accuracy jobs failed at import with ModuleNotFoundError: aiter.ops.triton.attention.mla_decode, raised from this branch's old triton_mla.py against the current aiter wheel, so a re-run of that merge ref would fail the same way. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Merged Why a merge instead of rerunning the failed jobs: in run 39385165211, the GLM-5.3-Flash, Kimi-K3 and Kimi-K3 DSpark accuracy jobs fail at @valarLip all five review threads now have replies, each pointing to the commit that addressed it (4124e7e; 0f805f0 also covers the retreated-joint follow-up). Verified (MI355X,
|
| # has nothing new to reject. | ||
| if not decision.state_from_hbm and extra == 0: | ||
| return decision | ||
| num_new_blocks = getattr(seq, "_admission_num_new_blocks", None) |
There was a problem hiding this comment.
When _admission_num_new_blocks is missing, _fit_joint returns the decision unchecked, so the joint is committed without the re-check that prevents the exact-fit crash in _fresh_block. It cannot be missing in today's call order (can_allocate sets it on every probe, record_allocation runs after it in the same pass), but that ordering is an implicit contract carried by an undeclared dynamic attribute on Sequence, and the guard fails open: a future caller that commits without probing first, or a probe path that returns before line 1263, would silently bring the crash back. It is also never cleared, so a stale value from an earlier admission attempt could be read if the order ever changes.
Could the demand travel with the probe instead, e.g. as a field on JointBoundaryDecision / the probe result, or as a declared field on Sequence reset after use? At minimum, fail closed here: skip the joint (self._no_joint(...)) or assert, rather than passing it through.
There was a problem hiding this comment.
Fixed in ead6826: the demand now travels with the probe, and the guard fails closed.
- New
AdmissionProbe(hit, walked chain,num_new_blocks).can_allocate(..., probe=...)resetsnum_new_blockstoNoneon every probe and sets it only once the fit check passes, so a refused re-probe cannot carry an earlier value. The_admission_num_new_blocksattribute onSequenceis gone. record_allocation(seq, probe)takes the probe. The scheduler's_probe_hbm_prefixreturns it and hands the same object torecord_allocation; a direct admission (record=True) builds one from_evaluate's locals._fit_joint(decision, probe): ifprobe.num_new_blocks is None, it returns_no_joint("no_probe_demand"). The check comes before every early return, including the tier-joint-without-private-tail shortcut, which never read the demand before. I chose a skip over an assert because skipping only falls back to the HBM hit. The new reason shows up injoint_skips.
Tests in tests/test_state_checkpoint.py::TestPagedCopyCheckpoint:
test_the_unchecked_exact_fit_joint_does_crash_allocate: on the exact-fit pool (now_exact_fit_boundary_image_pool), committing the joint without_fit_jointmakesallocateraiseNo PAGE unit available.test_a_joint_without_probe_demand_is_skipped_not_committed: on the same pool, a probe without demand skips the joint (joint_skips == {"no_probe_demand": 1}), andallocatesucceeds at the HBM hit.test_a_tier_joint_without_probe_demand_is_skipped_too: covers the shortcut path.test_a_refused_probe_keeps_no_demand_from_an_earlier_one: the stale-value case.
Restoring the old pass-through makes the two no_probe_demand tests fail.
| h = seq.cache_seed | ||
| for i in range(hbm_blocks): | ||
| token_ids = self._hash_block_tokens(seq, i) | ||
| h = self.compute_hash(token_ids, h) |
There was a problem hiding this comment.
This recomputes the prompt's hash chain from cache_seed block by block, while can_allocate / allocate use _prompt_hashes(seq), which may adopt the frontend's shipped chain (_adopt_shipped_chain). Two sources for the same chain: if they ever diverge (hash function, seed handling, block size), every lookup here misses and the retreat silently fails, falling back to token 0 -- exactly the recompute this function exists to avoid, with nothing in the logs to say why. It is also an extra O(hbm_blocks) hash pass on a failure path that can cover hundreds of thousands of tokens.
Reading the first hbm_blocks entries of self._prompt_hashes(seq) here (the chain can_allocate walked) would keep one source of truth and drop the rehash; the token_ids comparison on each block can stay as the collision check.
There was a problem hiding this comment.
Fixed in ead6826. retreat_joint_to_hbm now reads the HBM hit from self._prompt_hashes(seq), the chain can_allocate walked, which may be the adopted frontend chain. It goes through _block_hash(seq, prompt, i, h) the way _walk_prefix does: an entry the chain already holds is read rather than rehashed, so the extra hash pass is gone. The per-block token_ids comparison stays as the collision check. A miss now logs [JOINT-RETREAT] seq ...: HBM hit block i of n is gone at debug level, so a failed retreat says where it stopped.
Test: test_a_failed_unaligned_joint_resumes_at_the_hbm_hit is parametrized. In the shipped_chain case, the adopted chain is minted from a different seed, so it differs from what this engine would hash from cache_seed (the test asserts the first link differs). The retreat still lands on the HBM hit, with the shared tail block and the restore queued. Switching the loop back to compute_hash from cache_seed makes that case fail.
…its chain Review follow-ups on the unaligned joint (#2562): - The demand `can_allocate`'s fit passed now travels with the probe (`AdmissionProbe`: hit, walked chain, `num_new_blocks`) into `record_allocation(seq, probe)` instead of an undeclared, never-cleared `Sequence` attribute. `_fit_joint` fails closed: a probe without demand skips the joint (`no_probe_demand`) before any shortcut. Each probe resets the demand, so a refused re-probe cannot carry an earlier one. - `retreat_joint_to_hbm` reads the HBM hit off `_prompt_hashes(seq)`, the chain `can_allocate` walked (possibly the frontend's), instead of rehashing from `cache_seed`. The per-block token comparison stays as the collision check; a miss is logged at debug. Tests: the unchecked exact-fit joint crashes `allocate`; without probe demand the joint is skipped (exact-fit and tier-joint cases); a refused re-probe clears the demand; the retreat resumes at the HBM hit with a shipped chain that differs from the locally hashed one.
|
@valarLip @ZhangLirong-amd Both inline comments are addressed in ead6826 (
Retreat rehashing from The PR description now covers both. Validation:
Merge check: #2636 still conflicts only in Could you take another look? |
Resolve against #2562 (035a78a), which added `AdmissionProbe`: - block_manager `_evaluate`: keep main's positional `probe` out-param and the PR's keyword-only `reserved_blocks` / `reserved_state_slots` / `demand` / `observe`; after the fit passes, fill both `demand` and `probe.num_new_blocks`. - scheduler fill probe: `_can_detach_tier_lookup` now takes an `AdmissionProbe`, so pass `AdmissionProbe(cached_blocks)` instead of the old `(hit, hashes)` tuple.
Motivation
On the in-process LMCache path (
ATOM_KV_OFFLOAD=lmcache), a Qwen3.5 request can resume from the CPU tier only at a boundaryBwhere the KV for[0, B)and a GDN state checkpoint atBboth exist. The KV leg reads LMCache from the HBM prefix hit up toB. Two constraints decide where that read may start:LMCACHE_CHUNK_SIZE(256).The HBM hit sits on the 16-token hash-block grid. Agentic resumes land on the previous prompt's end, which is rarely a multiple of 256. With an unaligned hit there was no legal start: rounding down writes into a shared block, and rounding up leaves a gap the forward skips.
_joint_kv_boundaryrefused the joint load (hbm_off_chunk_grid). The request then recomputed everything above the HBM hit, even when LMCache held the KV and the state tier held a rung far above it.This change moves the shared/private split instead of the start:
_joint_kv_boundaryno longer refuses an unaligned HBM hit.claim_tokensis the compressed-hit chunk floor when that floor is already at or above the gated hit, so resident KV stays shared. Only a floor still below the hit steps down to the hit's own chunk floor.record_allocationchecks that floor before the joint is counted. The extra demand is the in-use blocks in[floor, hbm), added to the admission's existing demand, and the checkpoint about to be pinned is protected. If that does not fit, the joint is skipped (partial_chunk_no_room) and the request keeps its HBM hit. The skip is not also counted as a boundary. The admission's demand comes from the samecan_allocateprobe (AdmissionProbe.num_new_blocks, handed torecord_allocation(seq, probe)). A probe without it skips the joint (no_probe_demand) instead of committing it unchecked.allocatethen stops the shared claim at that floor. The partial chunk gets fresh private blocks after the restore is pinned.num_cached_tokensdrops to the floor. The KV leg starts atmax(hbm, claim_tokens)in the K3 connector, which is on the grid, and writes only private blocks up to the chunk covering the state rung.retreat_joint_to_hbmputs the HBM hit's shared blocks back and binds that hit's state into the slot the admission holds. It finds the hit on_prompt_hashes(seq), the chaincan_allocatewalked, which may be the frontend's, and compares tokens on each block. Only when the hit's blocks or state are gone does the request disown to token 0.Example from the new test (block 4 tokens, chunk 16): the HBM hit is 20 tokens and the state tier has a rung at 32. Before, the joint load was refused and the forward ran from 20. After,
[0, 16)stays shared,[16, 20)moves to a private block, LMCache loads[16, 32), the state loads at 32, and the forward runs from 32.Cost: the private refill is at most one chunk — 255 tokens and 15 blocks — and only when the compressed hit does not already reach a higher chunk boundary. A longer resident prefix stays shared and is not refetched. Aligned hits and models without per-request state are unchanged.
Command
Unit tests:
python -m pytest tests/test_state_checkpoint.py -q python -m pytest tests/test_block_manager.py tests/test_lmcache_*.py \ tests/test_prefill_prefix_vs_native.py tests/test_prefix_cache_accuracy.py -q188 and 716 passed on the previous revision. This revision adds the live-tail exact-fit case, the compressed-hit floor above the gated hit, and the
partial_chunk_no_roomfallback.tests/test_state_checkpoint.pyandtests/test_block_manager.pypass (276), as do the joint and unaligned cases intests/test_lmcache_offload_connector.py(18). The original case isTestPagedCopyCheckpoint::test_an_unaligned_hbm_hit_still_reaches_the_rung_above_it: the hit is gated to 5 blocks (20 tokens), the joint boundary is 8 blocks,claim_tokens == 16, block 3 is still the shared block, block 4 is a fresh block, andnum_cached_tokens == 16.Serving, InferenceX AgentX at concurrency 56. Model
/models/Qwen/Qwen3.5-397B-A17B-FP8, TP4, FP8 KV, LMCache CPU 253 GiB per rank:HIP_VISIBLE_DEVICES=0,1,2,3 AITER_LOG_LEVEL=WARNING PYTHONHASHSEED=0 \ ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 \ MAX_NUM_SEQS=112 GPU_MEM_UTIL=0.8 \ ATOM_KV_OFFLOAD=lmcache \ LMCACHE_LOCAL_CPU=True LMCACHE_NUMA_MODE=auto \ LMCACHE_MAX_LOCAL_CPU_SIZE=253 LMCACHE_CHUNK_SIZE=256 \ OFFLOAD_MIN_LOAD_TOKENS=8192 \ bash scripts/start_atom_server.sh /models/Qwen/Qwen3.5-397B-A17B-FP8 4 8000 \ --enable-prefix-caching --method mtp --num-speculative-tokens 3 --block-size 16 \ --cudagraph-capture-sizes "[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112]" aiperf profile --scenario inferencex-agentx-mvp \ --url http://127.0.0.1:8000 --endpoint /v1/chat/completions --endpoint-type chat --streaming \ --model /models/Qwen/Qwen3.5-397B-A17B-FP8 --tokenizer /models/Qwen/Qwen3.5-397B-A17B-FP8 \ --tokenizer-trust-remote-code --apply-chat-template \ --public-dataset semianalysis_cc_traces_weka_062126_256k --num-dataset-entries 393 \ --concurrency 56 --benchmark-duration 3600 --random-seed 42 \ --warmup-requests-per-lane 10 --warmup-grace-period 1800 --trace-idle-gap-cap-seconds 300 \ --trajectory-start-min-ratio 0.25 --trajectory-start-max-ratio 0.75 \ --failed-request-threshold 0.10 --use-server-token-count --no-gpu-telemetry \ --stats-interval 30 --slice-duration 1.0PYTHONHASHSEED=0is required, or cache keys will not match across workers. The state checkpoint interval is the default 8192 and the prefill chunk is 16384.Accuracy, same protocol as #2519 (TP8, FP8 KV, in-process LMCache, 32 GB CPU pool):
With
CHATunset,run_gsm8k_eval.shuseslocal-completions, 3-shot,num_concurrent=65.Performance
One 3600 s run each, with identical server and client settings. "Before" and "This PR" use LMCache's default rank-0 lookup. The third column runs this PR with all-rank lookup, a configuration change described below that is not part of this PR's code. The hit rate is over successful requests in the measured phase,
sum(cached_tokens) / sum(prompt_tokens), with warmup excluded.Joint-load counters at the end of each run (lifetime, including warmup):
hbm_off_chunk_gridis gone. Admissions it used to refuse now reach the rung search. Most find no CPU rung above the HBM hit (no_rung_above_hbm), and the rest become joint loads.partial_chunk_no_roomdid not fire in either run.A fixed-duration replay serves a different set and order of requests in each run, and run-to-run variation was not measured, so the 0.73-point difference between "Before" and "This PR" is not attributed to this change alone.
Joint-load failures against boundaries
Load-channel failures divided by
boundaries(lifetime, including warmup):On the rank-0 run, every KV retrieve except one started at token 0 (
required == lookup total). The exception, request 3643, transferred 31488 tokens of a 32768 lookup, so the prefix below the transfer was 1280 tokens. None of the 16 failures dropped a long HBM hit in order to extend a short distance past it. Ten of the 16 are the rank-0 lookup misses already described (three ranks retrieve 0, one retrieves the chunk). The all-rank run removes that pattern; its 6 remaining failures are state-leg failures, which still disown to token 0. That is 6.5% of boundaries, not a case of recomputing ~200k tokens to save one chunk.Re-run after the review fixes (
05e502613)Same rank-0 protocol as the "This PR" column above (TP4, MTP3, concurrency 56, 3600 s measured phase, warmup excluded). One run, on 2026-10-15.
05e502613The request mix differs across fixed-duration replays, as noted above, so the higher throughput and the 0.18-point hit-rate difference are not attributed to the review fixes alone. Nothing in this run is worse than the previous "This PR" column.
Re-run after merging main (
658aa3891)Same rank-0 protocol (TP4, MTP3, concurrency 56, 3600 s measured phase, warmup excluded), on 2026-10-20, after merging main. aiperf 0.12.0 ignores
--warmup-grace-periodon this scenario, so the warmup cap was--agentic-warmup-grace-period 1800. Warmup finished in 367 s with 0 errors. Figures are fromprofile_export_aiperf.json.05e502613658aa3891Lifetime joint counter at the end of the run, warmup included:
partial_chunk_no_roomdid not fire.state_offload_loads_failedstayed 0. The measured-phase GPU/CPU tier split is not in this export. The cumulative tier line at the end of profiling includes warmup (HBM 90.99%, CPU 1.98%) and is not comparable to the measured-phase row above.The request mix differs across fixed-duration replays, and this commit also contains main's draft-lookahead and midstep checkpoint changes, so the throughput and ITL differences are not attributed to the merge resolution.
All-rank lookup
LMCache starts a lookup server only on the workers listed in
lookup_server_worker_ids, and ATOM defaults that list to[0](atom/kv_transfer/offload/config.py). A lookup pins the matched CPU objects and refreshes their LRU entries only on the workers it queries. With rank 0 alone, ranks 1–3 evict their shards of the same chunk independently. A retrieve then finds a chunk missing on some ranks. The load needs every rank, so the whole load fails and the request recomputes from token 0.In the "This PR" run, 10 of the 16 load failures had three ranks retrieve 0 tokens and one rank retrieve everything. We matched the 3876 successful requests one-to-one to the SGLang run on the same trace (InferenceX run 36743632401). On 8 of the failed requests, SGLang served 1,296,000 cached tokens that ATOM recomputed.
The third column only adds:
state_offload_loads_failed=6).InvalidInferenceResultError, no content returned), against 1 connection reset in the "This PR" run. The root cause is not identified. The three checked so far had no CPU state load while they ran.This PR does not change the default. The lookup scope stays rank 0 unless the configuration above is set.
Acc
Full GSM8K, 1319 questions,
local-completions, 3-shot, concurrency 65,max_num_seqs=65, in-process LMCache, on this commit. The request phase took about 7 min.Both metrics are 0.003–0.005 lower than in #2519, within one standard error of either run.
Re-run on
05e502613, same protocol (TP8,local-completions, 3-shot, concurrency 65,max_num_seqs=65, in-process LMCache, 32 GB CPU pool). The request phase took about 6 min.05e502613flexible-extract matches the previous run exactly. strict-match is 0.0023 lower, inside one standard error of either run.
Re-run on
658aa3891, same protocol. The request phase took about 5 min.658aa389105e502613Both stay above the 0.85 threshold. flexible-extract is 0.0045 lower than
05e502613, inside one standard error of either run.